package io.kestra.plugin.core.storage;

import com.fasterxml.jackson.core.type.TypeReference;
import io.kestra.core.models.property.Property;
import io.kestra.core.serializers.JacksonMapper;
import io.swagger.v3.oas.annotations.media.Schema;
import lombok.*;
import lombok.experimental.SuperBuilder;
import org.apache.commons.io.IOUtils;
import io.kestra.core.models.annotations.Example;
import io.kestra.core.models.annotations.Plugin;
import io.kestra.core.models.annotations.PluginProperty;
import io.kestra.core.models.tasks.RunnableTask;
import io.kestra.core.models.tasks.Task;
import io.kestra.core.runners.RunContext;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.InputStream;
import java.net.URI;
import java.util.List;

import jakarta.validation.constraints.NotNull;

import static io.kestra.core.utils.Rethrow.throwConsumer;

@SuperBuilder
@ToString
@EqualsAndHashCode
@Getter
@NoArgsConstructor
@Schema(
    title = "Concat files from Kestra’s internal storage."
)
@Plugin(
    examples = {
        @Example(
            title = "Concat 2 files with a custom separator.",
            code = {
                "files: ",
                "  - \"kestra://long/url/file1.txt\"",
                "  - \"kestra://long/url/file2.txt\"",
                "separator: \"\\n\"",
            }
        ),
        @Example(
            title = "Concat files generated by an each task.",
            code = {
                "tasks:",
                "  - id: each",
                "    type: io.kestra.plugin.core.flow.ForEach",
                "    tasks:",
                "      - id: start_api_call",
                "        type: io.kestra.plugin.scripts.shell.Commands",
                "        commands:",
                "          - echo {{ taskrun.value }} > {{ temp.generated }}",
                "        files:",
                "          - generated",
                "    values: '[\"value1\", \"value2\", \"value3\"]'",
                "",
                "  - id: concat",
                "    type: io.kestra.plugin.core.storage.Concat",
                "    files:",
                "      - \"{{ outputs.start_api_call.value1.files.generated }}\"",
                "      - \"{{ outputs.start_api_call.value2.files.generated }}\"",
                "      - \"{{ outputs.start_api_call.value3.files.generated }}\"",
            },
            full = true
        ),
        @Example(
            title = "Concat a dynamic number of files.",
            code = {
                "tasks:",
                "  - id: echo",
                "    type: io.kestra.plugin.scripts.shell.Commands",
                "    commands:",
                "      - echo \"Hello John\" > {{ outputDirs.output }}/1.txt",
                "      - echo \"Hello Jane\" > {{ outputDirs.output }}/2.txt",
                "      - echo \"Hello Doe\" > {{ outputDirs.output }}/3.txt",
                "    outputDirs:",
                "      - output",
                "  - id: concat",
                "    type: io.kestra.plugin.core.storage.Concat",
                "    files: \"{{ outputs.echo.files | jq('.[]') }}\"",
            },
            full = true
        )
    },
    aliases = "io.kestra.core.tasks.storages.Concat"
)
public class Concat extends Task implements RunnableTask<Concat.Output> {
    @Schema(
        title = "List of files to be concatenated.",
        description = "Must be a `kestra://` storage URIs, can be a list of string or json string"
    )
    @PluginProperty(dynamic = true, internalStorageURI = true)
    @NotNull
    private Object files;

    @Schema(
        title = "The separator to used between files, default is no separator."
    )
    private Property<String> separator;

    @Schema(
        title = "The extension of the created file, default is .tmp."
    )
    @Builder.Default
    private Property<String> extension = Property.ofValue(".tmp");

    @SuppressWarnings("unchecked")
    @Override
    public Concat.Output run(RunContext runContext) throws Exception {
        File tempFile = runContext.workingDir().createTempFile(runContext.render(extension).as(String.class).orElseThrow()).toFile();
        try (FileOutputStream fileOutputStream = new FileOutputStream(tempFile)) {
            List<String> finalFiles;
            if (this.files instanceof List<?> listValue) {
                finalFiles = (List<String>) listValue;
            } else if (this.files instanceof String stringValue) {
                final TypeReference<List<String>> reference = new TypeReference<>() {};

                finalFiles = JacksonMapper.ofJson(false).readValue(
                    runContext.render(stringValue),
                    reference
                );
            } else {
                throw new Exception("Invalid `files` properties with type '" + this.files.getClass() + "'");
            }

            finalFiles.forEach(throwConsumer(s -> {
                URI from = new URI(runContext.render(s));
                try (InputStream inputStream = runContext.storage().getFile(from)) {
                    IOUtils.copyLarge(inputStream, fileOutputStream);
                }

                if (separator != null) {
                    IOUtils.copy(new ByteArrayInputStream(runContext.render(this.separator).as(String.class).orElseThrow().getBytes()), fileOutputStream);
                }
            }));
        }

        return Concat.Output.builder()
            .uri(runContext.storage().putFile(tempFile))
            .build();
    }

    @Builder
    @Getter
    public static class Output implements io.kestra.core.models.tasks.Output {
        @Schema(
            title = "The concatenated file URI."
        )
        private final URI uri;
    }
}
